import { GPU } from "gpu.js";
import { writeByte8 } from "../lib/utils";
import { filters } from "./filters";
const gpu = new GPU();

let byteWidth = 8;
const bpp = 4;

const filterSumNone = filterSumFactor(filterSumNoneGPUFunc);
const filterSumSub = filterSumFactor(filterSumSubGPUFunc);
const filterSumUp = filterSumFactor(filterSumUpGPUFunc);
const filterSumAvg = filterSumFactor(filterSumAvgGPUFunc);
const filterSumPaeth = filterSumFactor(filterSumPaethGPUFunc);
const filterSums = [
  filterSumNone,
  filterSumSub,
  filterSumUp,
  filterSumAvg,
  filterSumPaeth,
];

export function filterGPU(
  imageData: ArrayBuffer,
  width: number,
  height: number
) {
  byteWidth = width * bpp;
  let rawPos = 0,
    pxPos = 0;
  const rawData = new ArrayBuffer((byteWidth + 1) * height);
  const matrix: Uint8Array[] = [];
  //将图像数据转为矩阵
  for (let y = 0; y < height; y++) {
    const pointer = y * byteWidth;
    const row = imageData.slice(pointer, pointer + byteWidth);
    matrix.push(new Uint8Array(row));
  }
  const minArr = Array.from({ length: height }, () => Infinity);
  const indexArr = Array.from({ length: height }, () => 0);
  for (let i = 1; i < filterSums.length; i++) {
    const result = filterSums[i](matrix as any, height);
    minArr.forEach((item, index) => {
      if (result[index] < item) indexArr[index] = i;
    });
  }
  for (let y = 0; y < height; y++) {
    const filterSelect = indexArr[y];
    // console.log(`第${y}行使用${filterSelect}`);
    writeByte8(rawData, [filterSelect], rawPos);
    rawPos++;
    filters[filterSelect](imageData, pxPos, rawData, rawPos, byteWidth);
    rawPos += byteWidth;
    pxPos += byteWidth;
  }
  return rawData;
}

function filterSumFactor(
  type: (pxData: number[][], byteWidth: number) => number
) {
  return function (imageData: number[][], height: number) {
    const kernel = gpu.createKernel(type).setOutput([1, height]);
    const data = kernel(imageData, byteWidth) as number[][];
    return data.map((item) => item[0]);
  };
}

function filterSumNoneGPUFunc(pxData: number[][], byteWidth: number) {
  let sum = 0;
  for (let i = 0; i < byteWidth; i++) {
    //@ts-ignore
    sum += pxData[this.thread.y][i];
  }
  return sum;
}

function filterSumSubGPUFunc(pxData: number[][], byteWidth: number) {
  let sum = 0;
  for (let i = 0; i < byteWidth; i++) {
    let left = 0;
    if (i >= 4) {
      //@ts-ignore
      left = pxData[this.thread.y][i - 4];
    }
    //@ts-ignore
    sum += pxData[this.thread.y][i] - left;
  }
  return sum;
}

function filterSumUpGPUFunc(pxData: number[][], byteWidth: number) {
  let sum = 0;
  for (let i = 0; i < byteWidth; i++) {
    let up = 0;
    //@ts-ignore
    if (this.thread.y > 1) {
      //@ts-ignore
      up = pxData[this.thread.y - 1][i];
    }
    //@ts-ignore
    sum += pxData[this.thread.y][i] - up;
  }
  return sum;
}

function filterSumAvgGPUFunc(pxData: number[][], byteWidth: number) {
  let sum = 0;
  for (let i = 0; i < byteWidth; i++) {
    let left = 0;
    let up = 0;
    if (i >= 4) {
      //@ts-ignore
      left = pxData[this.thread.y][i - 4];
    }
    //@ts-ignore
    if (this.thread.y > 1) {
      //@ts-ignore
      up = pxData[this.thread.y - 1][i];
    }
    //@ts-ignore
    sum += pxData[this.thread.y][this.thread.x] - ((left + up) >> 1);
  }
  return sum;
}

function filterSumPaethGPUFunc(pxData: number[][], byteWidth: number) {
  let sum = 0;
  for (let i = 0; i < byteWidth; i++) {
    let left = 0;
    let up = 0;
    let upleft = 0;
    let target = 0;
    if (i >= 4) {
      //@ts-ignore
      left = pxData[this.thread.y][i - 4];
    }
    //@ts-ignore
    if (this.thread.y > 1) {
      //@ts-ignore
      up = pxData[this.thread.y - 1][i];
    }
    //@ts-ignore
    if (i >= 4 && this.thread.y > 1) {
      //@ts-ignore
      upleft = pxData[this.thread.y - 1][i - 4];
    }
    const paeth = left + up - upleft;
    const pLeft = Math.abs(paeth - left);
    const pAbove = Math.abs(paeth - up);
    const pUpLeft = Math.abs(paeth - upleft);

    if (pLeft <= pAbove && pLeft <= pUpLeft) {
      target = left;
    } else if (pAbove <= pUpLeft) {
      target = up;
    } else {
      target = upleft;
    }
    //@ts-ignore
    sum += pxData[this.thread.y][this.thread.x] - target;
  }
  return sum;
}
